from transformers import pipeline, AutoTokenizer
from scipy.special import logit

model = 'projecte-aina/roberta-base-ca-cased-sts'
tokenizer = AutoTokenizer.from_pretrained(model)
pipe = pipeline('text-classification', model=model, tokenizer=tokenizer)


def prepare(sentence_pairs):
    sentence_pairs_prep = []
    for s1, s2 in sentence_pairs:
        sentence_pairs_prep.append(
            f"{tokenizer.cls_token} {s1}{tokenizer.sep_token}{tokenizer.sep_token} {s2}{tokenizer.sep_token}")
    return sentence_pairs_prep


def get_score(sen1, sen2):
    sentence_pairs = [(sen1, sen2)]
    prediction = pipe(prepare(sentence_pairs), add_special_tokens=False)
    print(logit(prediction[0]['score']) / 5)


if __name__ == '__main__':
    get_score("El llibre va caure per la finestra.", "El llibre va sortir volant.")
